//! 聚合算子实现
//!
//! 实现各种聚合操作，包括GroupBy、Count、Sum、Avg等

use super::*;
use std::collections::HashMap;
use std::time::Instant;

/// 聚合算子
/// 实现分组聚合操作
#[derive(Debug)]
pub struct AggregationOperator {
    /// 输入算子
    input: Box<dyn Operator>,
    /// 分组列索引
    group_by_columns: Vec<usize>,
    /// 聚合表达式
    aggregate_expressions: Vec<AggregateExpression>,
    /// 输出模式
    schema: ExecutionSchema,
    /// 聚合结果
    aggregation_results: Option<Vec<ExecutionTuple>>,
    /// 当前结果索引
    current_result_index: usize,
    /// 统计信息
    statistics: OperatorStatistics,
    /// 是否已打开
    is_open: bool,
}

/// 排序算子
/// 实现ORDER BY操作
#[derive(Debug)]
pub struct SortOperator {
    /// 输入算子
    input: Box<dyn Operator>,
    /// 排序键
    sort_keys: Vec<SortKey>,
    /// 排序结果
    sorted_results: Option<Vec<ExecutionTuple>>,
    /// 当前结果索引
    current_result_index: usize,
    /// 输出模式（与输入相同）
    schema: ExecutionSchema,
    /// 统计信息
    statistics: OperatorStatistics,
    /// 是否已打开
    is_open: bool,
}

/// 限制算子
/// 实现LIMIT和OFFSET操作
#[derive(Debug)]
pub struct LimitOperator {
    /// 输入算子
    input: Box<dyn Operator>,
    /// 限制数量
    limit: Option<usize>,
    /// 偏移量
    offset: usize,
    /// 当前处理的元组数量
    current_count: usize,
    /// 输出模式（与输入相同）
    schema: ExecutionSchema,
    /// 统计信息
    statistics: OperatorStatistics,
    /// 是否已打开
    is_open: bool,
}

/// 聚合表达式
#[derive(Debug, Clone)]
pub struct AggregateExpression {
    /// 聚合函数类型
    pub function: AggregateFunction,
    /// 输入列索引
    pub input_column: usize,
    /// 输出列名
    pub output_name: String,
}

/// 聚合函数类型
#[derive(Debug, Clone, PartialEq)]
pub enum AggregateFunction {
    Count,
    Sum,
    Avg,
    Min,
    Max,
    CountDistinct,
}

/// 排序键
#[derive(Debug, Clone)]
pub struct SortKey {
    /// 列索引
    pub column_index: usize,
    /// 排序方向
    pub direction: SortDirection,
}

/// 排序方向
#[derive(Debug, Clone, PartialEq)]
pub enum SortDirection {
    Ascending,
    Descending,
}

/// 聚合状态
#[derive(Debug, Clone)]
enum AggregateState {
    Count(u64),
    Sum(Value),
    Avg { sum: Value, count: u64 },
    Min(Value),
    Max(Value),
    CountDistinct(std::collections::HashSet<Vec<u8>>),
}

impl AggregationOperator {
    /// 创建新的聚合算子
    pub fn new(
        input: Box<dyn Operator>,
        group_by_columns: Vec<usize>,
        aggregate_expressions: Vec<AggregateExpression>,
    ) -> Result<Self> {
        let input_schema = input.schema();
        
        // 构建输出模式
        let mut output_columns = Vec::new();
        
        // 添加分组列
        for &column_index in &group_by_columns {
            if let Some(column) = input_schema.get_column(column_index) {
                output_columns.push(column.clone());
            } else {
                return Err(Error::query_execution(format!("分组列索引 {} 超出范围", column_index)));
            }
        }
        
        // 添加聚合列
        for agg_expr in &aggregate_expressions {
            let data_type = match agg_expr.function {
                AggregateFunction::Count | AggregateFunction::CountDistinct => DataType::BigInt,
                AggregateFunction::Sum | AggregateFunction::Avg => {
                    // 根据输入列类型确定输出类型
                    if let Some(input_column) = input_schema.get_column(agg_expr.input_column) {
                        match input_column.data_type {
                            DataType::Integer => DataType::BigInt,
                            DataType::Float => DataType::Double,
                            _ => DataType::Double,
                        }
                    } else {
                        DataType::Double
                    }
                }
                AggregateFunction::Min | AggregateFunction::Max => {
                    // 保持输入类型
                    if let Some(input_column) = input_schema.get_column(agg_expr.input_column) {
                        input_column.data_type.clone()
                    } else {
                        return Err(Error::query_execution(format!("聚合列索引 {} 超出范围", agg_expr.input_column)));
                    }
                }
            };
            
            output_columns.push(ColumnInfo::new(
                agg_expr.output_name.clone(),
                data_type,
                false, // 聚合结果通常不为NULL
            ));
        }
        
        let schema = ExecutionSchema::new(output_columns);
        
        Ok(Self {
            input,
            group_by_columns,
            aggregate_expressions,
            schema,
            aggregation_results: None,
            current_result_index: 0,
            statistics: OperatorStatistics::default(),
            is_open: false,
        })
    }
}

impl Operator for AggregationOperator {
    fn open(&mut self, context: &mut ExecutionContext) -> Result<()> {
        if self.is_open {
            return Ok(());
        }
        
        let start_time = Instant::now();
        
        self.input.open(context)?;
        
        // 执行聚合计算
        self.compute_aggregation(context)?;
        
        self.is_open = true;
        self.statistics.execution_time_us += start_time.elapsed().as_micros() as u64;
        
        tracing::debug!("聚合算子已打开");
        Ok(())
    }
    
    fn next(&mut self, _context: &mut ExecutionContext) -> Result<Option<ExecutionTuple>> {
        if !self.is_open {
            return Err(Error::query_execution("算子未打开".to_string()));
        }
        
        let start_time = Instant::now();
        
        if let Some(ref results) = self.aggregation_results {
            if self.current_result_index < results.len() {
                let result = results[self.current_result_index].clone();
                self.current_result_index += 1;
                
                self.statistics.tuples_output += 1;
                self.statistics.execution_time_us += start_time.elapsed().as_micros() as u64;
                
                Ok(Some(result))
            } else {
                self.statistics.execution_time_us += start_time.elapsed().as_micros() as u64;
                Ok(None)
            }
        } else {
            self.statistics.execution_time_us += start_time.elapsed().as_micros() as u64;
            Ok(None)
        }
    }
    
    fn close(&mut self, context: &mut ExecutionContext) -> Result<()> {
        self.input.close(context)?;
        self.aggregation_results = None;
        self.current_result_index = 0;
        self.is_open = false;
        
        tracing::debug!("聚合算子已关闭");
        Ok(())
    }
    
    fn schema(&self) -> &ExecutionSchema {
        &self.schema
    }
    
    fn statistics(&self) -> OperatorStatistics {
        let mut stats = self.statistics.clone();
        let input_stats = self.input.statistics();
        stats.disk_io_count += input_stats.disk_io_count;
        stats.memory_usage_bytes += input_stats.memory_usage_bytes;
        stats
    }
    
    fn reset(&mut self) -> Result<()> {
        self.input.reset()?;
        self.aggregation_results = None;
        self.current_result_index = 0;
        self.statistics = OperatorStatistics::default();
        self.is_open = false;
        Ok(())
    }
}

impl AggregationOperator {
    /// 计算聚合结果
    fn compute_aggregation(&mut self, context: &mut ExecutionContext) -> Result<()> {
        let mut groups: HashMap<Vec<u8>, HashMap<usize, AggregateState>> = HashMap::new();
        
        // 遍历输入元组
        while let Some(tuple) = self.input.next(context)? {
            self.statistics.tuples_processed += 1;
            
            // 提取分组键
            let group_key = self.extract_group_key(&tuple)?;
            
            // 获取或创建分组状态
            let group_aggregates = groups.entry(group_key).or_insert_with(HashMap::new);
            
            // 更新聚合状态
            for (agg_index, agg_expr) in self.aggregate_expressions.iter().enumerate() {
                let input_value = tuple.get_value(agg_expr.input_column)?;
                
                let aggregate_state = group_aggregates.entry(agg_index).or_insert_with(|| {
                    self.create_initial_aggregate_state(&agg_expr.function)
                });
                
                self.update_aggregate_state(aggregate_state, &agg_expr.function, input_value)?;
            }
        }
        
        // 生成最终结果
        let mut results = Vec::new();
        for (group_key, group_aggregates) in groups {
            let mut result_values = Vec::new();
            
            // 添加分组列值
            let group_values = self.decode_group_key(&group_key)?;
            result_values.extend(group_values);
            
            // 添加聚合结果
            for agg_index in 0..self.aggregate_expressions.len() {
                if let Some(aggregate_state) = group_aggregates.get(&agg_index) {
                    let final_value = self.finalize_aggregate_state(aggregate_state)?;
                    result_values.push(final_value);
                } else {
                    result_values.push(Value::Null);
                }
            }
            
            results.push(ExecutionTuple::new(result_values));
        }
        
        self.aggregation_results = Some(results);
        Ok(())
    }
    
    /// 提取分组键
    fn extract_group_key(&self, tuple: &ExecutionTuple) -> Result<Vec<u8>> {
        let mut key = Vec::new();
        
        for &column_index in &self.group_by_columns {
            let value = tuple.get_value(column_index)?;
            let value_bytes = self.serialize_value(value)?;
            key.extend(value_bytes);
            key.push(0); // 分隔符
        }
        
        Ok(key)
    }
    
    /// 序列化值为字节数组
    fn serialize_value(&self, value: &Value) -> Result<Vec<u8>> {
        match value {
            Value::Integer(i) => Ok(i.to_le_bytes().to_vec()),
            Value::Real(f) => Ok(f.to_le_bytes().to_vec()),
            Value::Text(s) => Ok(s.as_bytes().to_vec()),
            Value::Boolean(b) => Ok(vec![if *b { 1 } else { 0 }]),
            Value::Null => Ok(vec![255]), // NULL的特殊标记
        }
    }
    
    /// 解码分组键
    fn decode_group_key(&self, _key: &[u8]) -> Result<Vec<Value>> {
        // 简化实现：返回空值
        // 实际实现需要根据分组列类型正确解码
        Ok(vec![Value::String("group".to_string()); self.group_by_columns.len()])
    }
    
    /// 创建初始聚合状态
    fn create_initial_aggregate_state(&self, function: &AggregateFunction) -> AggregateState {
        match function {
            AggregateFunction::Count => AggregateState::Count(0),
            AggregateFunction::Sum => AggregateState::Sum(Value::Integer(0)),
            AggregateFunction::Avg => AggregateState::Avg { sum: Value::Integer(0), count: 0 },
            AggregateFunction::Min => AggregateState::Min(Value::Null),
            AggregateFunction::Max => AggregateState::Max(Value::Null),
            AggregateFunction::CountDistinct => AggregateState::CountDistinct(std::collections::HashSet::new()),
        }
    }
    
    /// 更新聚合状态
    fn update_aggregate_state(
        &self,
        state: &mut AggregateState,
        function: &AggregateFunction,
        value: &Value,
    ) -> Result<()> {
        match (state, function) {
            (AggregateState::Count(count), AggregateFunction::Count) => {
                if !matches!(value, Value::Null) {
                    *count += 1;
                }
            }
            (AggregateState::Sum(sum), AggregateFunction::Sum) => {
                if !matches!(value, Value::Null) {
                    *sum = self.add_values(sum, value)?;
                }
            }
            (AggregateState::Avg { sum, count }, AggregateFunction::Avg) => {
                if !matches!(value, Value::Null) {
                    *sum = self.add_values(sum, value)?;
                    *count += 1;
                }
            }
            (AggregateState::Min(min_val), AggregateFunction::Min) => {
                if !matches!(value, Value::Null) {
                    if matches!(min_val, Value::Null) || self.compare_values(value, min_val)? < 0 {
                        *min_val = value.clone();
                    }
                }
            }
            (AggregateState::Max(max_val), AggregateFunction::Max) => {
                if !matches!(value, Value::Null) {
                    if matches!(max_val, Value::Null) || self.compare_values(value, max_val)? > 0 {
                        *max_val = value.clone();
                    }
                }
            }
            (AggregateState::CountDistinct(set), AggregateFunction::CountDistinct) => {
                if !matches!(value, Value::Null) {
                    let value_bytes = self.serialize_value(value)?;
                    set.insert(value_bytes);
                }
            }
            _ => return Err(Error::query_execution("聚合状态与函数不匹配".to_string())),
        }
        
        Ok(())
    }
    
    /// 完成聚合状态计算
    fn finalize_aggregate_state(&self, state: &AggregateState) -> Result<Value> {
        match state {
            AggregateState::Count(count) => Ok(Value::Integer(*count as i64)),
            AggregateState::Sum(sum) => Ok(sum.clone()),
            AggregateState::Avg { sum, count } => {
                if *count == 0 {
                    Ok(Value::Null)
                } else {
                    self.divide_values(sum, &Value::Integer(*count as i64))
                }
            }
            AggregateState::Min(min_val) => Ok(min_val.clone()),
            AggregateState::Max(max_val) => Ok(max_val.clone()),
            AggregateState::CountDistinct(set) => Ok(Value::Integer(set.len() as i64)),
        }
    }
    
    /// 值相加
    fn add_values(&self, a: &Value, b: &Value) -> Result<Value> {
        match (a, b) {
            (Value::Integer(x), Value::Integer(y)) => Ok(Value::Integer(x + y)),
            (Value::Real(x), Value::Real(y)) => Ok(Value::Real(x + y)),
            (Value::Integer(x), Value::Real(y)) => Ok(Value::Real(*x as f64 + y)),
            (Value::Real(x), Value::Integer(y)) => Ok(Value::Real(x + *y as f64)),
            _ => Err(Error::query_execution("不支持的值相加操作".to_string())),
        }
    }
    
    /// 值相除
    fn divide_values(&self, a: &Value, b: &Value) -> Result<Value> {
        match (a, b) {
            (Value::Integer(x), Value::Integer(y)) => {
                if *y == 0 {
                    Err(Error::query_execution("除零错误".to_string()))
                } else {
                    Ok(Value::Real(*x as f64 / *y as f64))
                }
            }
            (Value::Real(x), Value::Real(y)) => {
                if *y == 0.0 {
                    Err(Error::query_execution("除零错误".to_string()))
                } else {
                    Ok(Value::Real(x / y))
                }
            }
            _ => Err(Error::query_execution("不支持的值相除操作".to_string())),
        }
    }
    
    /// 比较值
    fn compare_values(&self, a: &Value, b: &Value) -> Result<i32> {
        use std::cmp::Ordering;
        
        let ordering = match (a, b) {
            (Value::Integer(x), Value::Integer(y)) => x.cmp(y),
            (Value::Real(x), Value::Real(y)) => x.partial_cmp(y).unwrap_or(Ordering::Equal),
            (Value::Text(x), Value::Text(y)) => x.cmp(y),
            _ => return Err(Error::query_execution("不支持的值比较操作".to_string())),
        };
        
        Ok(match ordering {
            Ordering::Less => -1,
            Ordering::Equal => 0,
            Ordering::Greater => 1,
        })
    }
}

impl SortOperator {
    /// 创建新的排序算子
    pub fn new(input: Box<dyn Operator>, sort_keys: Vec<SortKey>) -> Self {
        let schema = input.schema().clone();

        Self {
            input,
            sort_keys,
            sorted_results: None,
            current_result_index: 0,
            schema,
            statistics: OperatorStatistics::default(),
            is_open: false,
        }
    }
}

impl Operator for SortOperator {
    fn open(&mut self, context: &mut ExecutionContext) -> Result<()> {
        if self.is_open {
            return Ok(());
        }

        let start_time = Instant::now();

        self.input.open(context)?;

        // 收集所有输入元组并排序
        self.sort_input_tuples(context)?;

        self.is_open = true;
        self.statistics.execution_time_us += start_time.elapsed().as_micros() as u64;

        tracing::debug!("排序算子已打开");
        Ok(())
    }

    fn next(&mut self, _context: &mut ExecutionContext) -> Result<Option<ExecutionTuple>> {
        if !self.is_open {
            return Err(Error::query_execution("算子未打开".to_string()));
        }

        let start_time = Instant::now();

        if let Some(ref results) = self.sorted_results {
            if self.current_result_index < results.len() {
                let result = results[self.current_result_index].clone();
                self.current_result_index += 1;

                self.statistics.tuples_output += 1;
                self.statistics.execution_time_us += start_time.elapsed().as_micros() as u64;

                Ok(Some(result))
            } else {
                self.statistics.execution_time_us += start_time.elapsed().as_micros() as u64;
                Ok(None)
            }
        } else {
            self.statistics.execution_time_us += start_time.elapsed().as_micros() as u64;
            Ok(None)
        }
    }

    fn close(&mut self, context: &mut ExecutionContext) -> Result<()> {
        self.input.close(context)?;
        self.sorted_results = None;
        self.current_result_index = 0;
        self.is_open = false;

        tracing::debug!("排序算子已关闭");
        Ok(())
    }

    fn schema(&self) -> &ExecutionSchema {
        &self.schema
    }

    fn statistics(&self) -> OperatorStatistics {
        let mut stats = self.statistics.clone();
        let input_stats = self.input.statistics();
        stats.disk_io_count += input_stats.disk_io_count;
        stats.memory_usage_bytes += input_stats.memory_usage_bytes;
        stats
    }

    fn reset(&mut self) -> Result<()> {
        self.input.reset()?;
        self.sorted_results = None;
        self.current_result_index = 0;
        self.statistics = OperatorStatistics::default();
        self.is_open = false;
        Ok(())
    }
}

impl SortOperator {
    /// 排序输入元组
    fn sort_input_tuples(&mut self, context: &mut ExecutionContext) -> Result<()> {
        let mut tuples = Vec::new();

        // 收集所有输入元组
        while let Some(tuple) = self.input.next(context)? {
            self.statistics.tuples_processed += 1;
            tuples.push(tuple);
        }

        // 估算内存使用量
        let estimated_memory = tuples.len() * 64; // 简化估算
        context.allocate_memory(estimated_memory)?;
        self.statistics.memory_usage_bytes += estimated_memory as u64;

        // 排序元组
        tuples.sort_by(|a, b| self.compare_tuples(a, b));

        self.sorted_results = Some(tuples);
        Ok(())
    }

    /// 比较两个元组
    fn compare_tuples(&self, a: &ExecutionTuple, b: &ExecutionTuple) -> std::cmp::Ordering {
        use std::cmp::Ordering;

        for sort_key in &self.sort_keys {
            let a_value = a.get_value(sort_key.column_index);
            let b_value = b.get_value(sort_key.column_index);

            let ordering = match (a_value, b_value) {
                (Ok(a_val), Ok(b_val)) => self.compare_values_for_sort(a_val, b_val),
                (Err(_), Ok(_)) => Ordering::Greater, // 错误值排在后面
                (Ok(_), Err(_)) => Ordering::Less,
                (Err(_), Err(_)) => Ordering::Equal,
            };

            let final_ordering = match sort_key.direction {
                SortDirection::Ascending => ordering,
                SortDirection::Descending => ordering.reverse(),
            };

            if final_ordering != Ordering::Equal {
                return final_ordering;
            }
        }

        Ordering::Equal
    }

    /// 比较值用于排序
    fn compare_values_for_sort(&self, a: &Value, b: &Value) -> std::cmp::Ordering {
        use std::cmp::Ordering;

        match (a, b) {
            (Value::Null, Value::Null) => Ordering::Equal,
            (Value::Null, _) => Ordering::Less, // NULL值排在前面
            (_, Value::Null) => Ordering::Greater,
            (Value::Integer(x), Value::Integer(y)) => x.cmp(y),
            (Value::Real(x), Value::Real(y)) => x.partial_cmp(y).unwrap_or(Ordering::Equal),
            (Value::Text(x), Value::Text(y)) => x.cmp(y),
            (Value::Boolean(x), Value::Boolean(y)) => x.cmp(y),
            _ => Ordering::Equal, // 不同类型默认相等
        }
    }
}

impl LimitOperator {
    /// 创建新的限制算子
    pub fn new(input: Box<dyn Operator>, limit: Option<usize>, offset: usize) -> Self {
        let schema = input.schema().clone();

        Self {
            input,
            limit,
            offset,
            current_count: 0,
            schema,
            statistics: OperatorStatistics::default(),
            is_open: false,
        }
    }
}

impl Operator for LimitOperator {
    fn open(&mut self, context: &mut ExecutionContext) -> Result<()> {
        if self.is_open {
            return Ok(());
        }

        self.input.open(context)?;
        self.is_open = true;

        tracing::debug!("限制算子已打开");
        Ok(())
    }

    fn next(&mut self, context: &mut ExecutionContext) -> Result<Option<ExecutionTuple>> {
        if !self.is_open {
            return Err(Error::query_execution("算子未打开".to_string()));
        }

        let start_time = Instant::now();

        // 检查是否已达到限制
        if let Some(limit) = self.limit {
            if self.current_count >= self.offset + limit {
                self.statistics.execution_time_us += start_time.elapsed().as_micros() as u64;
                return Ok(None);
            }
        }

        // 跳过偏移量
        while self.current_count < self.offset {
            if let Some(_) = self.input.next(context)? {
                self.current_count += 1;
                self.statistics.tuples_processed += 1;
            } else {
                self.statistics.execution_time_us += start_time.elapsed().as_micros() as u64;
                return Ok(None);
            }
        }

        // 返回限制范围内的元组
        if let Some(tuple) = self.input.next(context)? {
            self.current_count += 1;
            self.statistics.tuples_processed += 1;
            self.statistics.tuples_output += 1;
            self.statistics.execution_time_us += start_time.elapsed().as_micros() as u64;

            Ok(Some(tuple))
        } else {
            self.statistics.execution_time_us += start_time.elapsed().as_micros() as u64;
            Ok(None)
        }
    }

    fn close(&mut self, context: &mut ExecutionContext) -> Result<()> {
        self.input.close(context)?;
        self.current_count = 0;
        self.is_open = false;

        tracing::debug!("限制算子已关闭");
        Ok(())
    }

    fn schema(&self) -> &ExecutionSchema {
        &self.schema
    }

    fn statistics(&self) -> OperatorStatistics {
        let mut stats = self.statistics.clone();
        let input_stats = self.input.statistics();
        stats.disk_io_count += input_stats.disk_io_count;
        stats.memory_usage_bytes += input_stats.memory_usage_bytes;
        stats
    }

    fn reset(&mut self) -> Result<()> {
        self.input.reset()?;
        self.current_count = 0;
        self.statistics = OperatorStatistics::default();
        self.is_open = false;
        Ok(())
    }
}
